"""
获取贝叶斯日报的数据，股票，期货，申万指数；
"""
import pandas as pd
from sqlalchemy import create_engine
from tools.utils import load_config
from tools.utils import loc_zq_collection


def load_engine(cfg):
    name = cfg['mysql']['user']
    pwd = cfg['mysql']['password']
    ip = cfg['mysql']['ip']
    engine = create_engine(f"mysql+pymysql://{name}:{pwd}@{ip}:3306/stock?charset=utf8&use_unicode=1")
    return engine


def load_stock_all(cfg):
    eg = load_engine(cfg)
    stock_names = pd.read_sql(f'show tables', eg).iloc[:, 0].to_list()
    c = {}
    for stock_name in stock_names:
        s = stock_name.replace('.', '_')
        try:
            data = pd.read_sql(f"select * from {s}", eg)
            if len(data) > 1000:
                c[stock_name] = data['close']
        except:
            print('miss', s)
    return c


cfg_ = load_config()


def load_sw_all():
    meta_li = list(loc_zq_collection('wind_wsd', 'metadata').find({}))
    mp = {v['__id']: v['alias'] for v in meta_li}
    c = {}
    for dic in list(loc_zq_collection('wind_wsd', 'data').find()):
        if len(dic['index']) > 1000:
            dt = pd.DataFrame(dic['data'], index=dic['index'], columns=dic['columns'])
            c[dic['columns'][0]] = dt
            dt.columns = [mp[dt.columns[0]]]
    return c


data = load_sw_all()
